"""
Tensorflow 2.x
"""
import tensorflow as tf
import numpy as np
from python_ai.common.xcommon import *

M, N = 100, 10
SHUFFLE_BUFFER = 2000
BATCH_SIZE = 32
x = np.random.normal(0, 1, [M, N])
y = np.random.choice(np.arange(N, dtype=np.int64), size=M, p=np.full(N, 1 / N))
print_numpy_ndarray_info(x, 'x')
print_numpy_ndarray_info(y, 'y')

ds = tf.data.Dataset.from_tensor_slices((x, y))\
    .shuffle(SHUFFLE_BUFFER)\
    .batch(batch_size=BATCH_SIZE, drop_remainder=True)\
    .prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

for i, (bx, by) in enumerate(ds):
    sep(i)
    bx = bx.numpy()
    by = by.numpy()
    print_numpy_ndarray_info(bx, 'bx')
    print_numpy_ndarray_info(by, 'by')
